using DataFrames
using QuadGK
using Distributions
using JLD
using Optim
using ForwardDiff
using CSV
using LinearAlgebra
using SpecialFunctions
using Random
using DelimitedFiles
using SparseArrays


clearconsole()

#-------------------------------------------------------------
# include Functions
#-------------------------------------------------------------
#cd("$(pwd())/Dropbox/Heterogeneity/Software/KS_Simulation/")
readDir = "$(pwd())/Functions/"
# include(readDir *"vech.jl");
#include(readDir *"VAR_Procedures.jl");
include(readDir *"VAR_MoreLags_Procedures.jl");
include(readDir *"Loaddata.jl");

#-------------------------------------------------------------
# choose specification files
#-------------------------------------------------------------
nVARSpec    = "1"
nVARMCMCSpec= "1"
specDir   = "$(pwd())/SpecFiles/"
include(specDir * "/AggVARspec" * nVARSpec * ".jl")
include(specDir * "/AggVARMCMCspec" * nVARMCMCSpec * ".jl")

#-------------------------------------------------------------
# load aggregate data
#-------------------------------------------------------------
dataDir  = "$(pwd())/data/"
agg_data, period_agg, ~ = loadaggdata(SampleStart,SampleEnd)
n_agg    = size(agg_data)[2]
data_adj = agg_data[Tdrop-nlags+1:end,:]

n = n_agg
p = nlags

# OLS estimation
YY, XX, PHIols, Sols, SIGMAols = data_to_YY(data_adj, nlags, nlags)

#-------------------------------------------------------------
# Generate prior
#-------------------------------------------------------------
igammadf = size(YY)[2]+sigdf
prior    = prior_ACP_stru(YY, SIGMAols, n_agg, igammadf, nlags, lambda)

#-------------------------------------------------------------
# Run Posterior Sampler
#-------------------------------------------------------------

println("")
println("Bayesian estimation of VAR: Gibbs Sampling... ")
println("")

store_alp, store_beta, store_Sig = sample_ThetaSig(YY,XX,PHIols,Sols,SIGMAols,nlags,prior,nsim)
store_Btilde, store_Sigpdraw     = getReducedForm(store_alp,store_beta,store_Sig)

# Initialize matrices for collecting draws from posterior
store_Bpdraw             = zeros(nsim,n*p,n)
store_Sigtrpdraw         = zeros(nsim,n,n)

for jj = 1:nsim

    Sigpdraw_j = Symmetric(store_Sigpdraw[jj,:,:])
    L = cholesky(Sigpdraw_j).L
    store_Sigtrpdraw[jj,:,:] = L

    Btilde = reshape(store_Btilde[jj,:], n*p,n)
    global store_Bpdraw[jj,:,:] = Btilde

end

store_Bpdraw             = store_Bpdraw[nburn+1:nsim,:,:];
store_Sigpdraw           = store_Sigpdraw[nburn+1:nsim,:,:];
store_Sigtrpdraw         = store_Sigtrpdraw[nburn+1:nsim,:,:];


#-------------------------------------------------------------
# Save draws
#-------------------------------------------------------------

sName    = "AggVAR" * nVARSpec * "_MCMC" * nVARMCMCSpec
savedir  = "$(pwd())/results/" * sName *"/";
try mkdir(savedir) catch; end
save(savedir * sName* "_PostDraws.jld", "Bpdraw", store_Bpdraw, "Sigpdraw", store_Sigpdraw, "Sigtrpdraw", store_Sigtrpdraw)

#-------------------------------------------------------------
# MISC ITEMS
#-------------------------------------------------------------

Bpmean       = dropdims(mean(store_Bpdraw,dims=1),dims=1)
Sigpmean     = dropdims(mean(store_Sigpdraw,dims=1),dims=1)
Sigtrpmean   = dropdims(mean(store_Sigtrpdraw,dims=1),dims=1)

save(savedir * sName* "_PostMeans.jld", "Bpmean", Bpmean, "Sigpmean", Sigpmean, "Sigtrpmean", Sigtrpmean)

# save OLS estiamtes and posterior means as CSV files
CSV.write(savedir * sName * "_PHIols.csv", DataFrame(PHIols,:auto))
CSV.write(savedir * sName * "_SIGMAols.csv", DataFrame(SIGMAols,:auto))
CSV.write(savedir * sName * "_Bpmean.csv", DataFrame(Bpmean,:auto))
CSV.write(savedir * sName * "_Sigpmean.csv", DataFrame(Sigpmean,:auto))
CSV.write(savedir * sName * "_Sigtrpmean.csv", DataFrame(Sigtrpmean,:auto))
